package basic.courseLearn07.problem;

import basic.courseLearn07.Edge;
import basic.courseLearn07.Graph;
import basic.courseLearn07.MyUnionFindSet;
import basic.courseLearn07.Node;
import org.junit.Test;

import java.util.*;

//Prim算法
public class PrimAlgorithm {

    class MinComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    //Prim算法，生成最小生成树
    public void prim(Graph graph,Node root) {
        // 1. 定义size，定义并查集，定义小根堆
        int size =1;
        MyUnionFindSet unionFindSet = new MyUnionFindSet();
        unionFindSet.init(new ArrayList<>(graph.nodes.values()));
        PriorityQueue<Edge> minHeap = new PriorityQueue<>(new MinComparator());
        // 2. root相邻的边加入小根堆
        System.out.print(root.value+"--> ");
        for (Edge edge : root.edges) {
            minHeap.add(edge);
        }
        // 3. 循环，直到size等于节点数
        while(size < graph.nodes.size()){
            Node tempNode=null;
            // 3.1循环选出满足条件的边
            while(!minHeap.isEmpty()){
                Edge edge = minHeap.poll();
                tempNode = edge.to;
                if (!unionFindSet.isSameSet(root,tempNode)){
                    break;
                }
            }
            // 3.2将选出的节点加入并查集
            unionFindSet.union(root,tempNode);
            System.out.print(tempNode.value+"--> ");
            // 3.3将选出的节点相邻的边加入小根堆
            for (Edge edge : tempNode.edges) {
                if (!unionFindSet.isSameSet(root,edge.to)){
                    minHeap.add(edge);
                }
            }
            size++;
        }
    }


    @Test
    public void test(){
        Graph graph = new Graph();
        Node node1 = new Node(1);
        Node node2 = new Node(2);
        Node node3 = new Node(3);
        Node node4 = new Node(4);
        Node node5 = new Node(5);
        HashMap<Integer, Node> map = new HashMap<>();


        Edge edge1 = new Edge(2, node1, node2);
        Edge edge2 = new Edge(1, node1, node4);
        Edge edge3  = new Edge(3, node1, node5);

        Edge edge4 = new Edge(2, node2, node1);
        Edge edge5 = new Edge(4, node2, node4);
        Edge edge6 = new Edge(4, node2, node5);
        Edge edge7 = new Edge(5, node2, node3);

        Edge edge8 = new Edge(5, node3, node2);
        Edge edge9 = new Edge(7, node3, node5);

        Edge edge10= new Edge(1, node4, node1);
        Edge edge11= new Edge(4, node4, node2);
        Edge edge12 = new Edge(3, node4, node5);

        Edge edge13 = new Edge(3, node5, node4);
        Edge edge14 = new Edge(3, node5, node1);
        Edge edge15 = new Edge(4, node5, node2);
        Edge edge16 = new Edge(7, node5, node3);

        node1.setEdges(new ArrayList<>(Arrays.asList(edge1, edge2, edge3)) );
        node2.setEdges(new ArrayList<>(Arrays.asList(edge4, edge5, edge6, edge7)) );
        node3.setEdges(new ArrayList<>(Arrays.asList(edge8, edge9)) );
        node4.setEdges(new ArrayList<>(Arrays.asList(edge10, edge11, edge12)) );
        node5.setEdges(new ArrayList<>(Arrays.asList(edge13, edge14, edge15, edge16)) );

        map.put(1, node1);
        map.put(2, node2);
        map.put(3, node3);
        map.put(4, node4);
        map.put(5, node5);
        graph.setNodes(map);

        HashSet<Edge> edges = new HashSet<>();
        edges.add(edge1);
        edges.add(edge2);
        edges.add(edge3);
        edges.add(edge4);
        edges.add(edge5);
        edges.add(edge6);
        edges.add(edge7);
        edges.add(edge8);
        edges.add(edge9);
        edges.add(edge10);
        edges.add(edge11);
        edges.add(edge12);
        edges.add(edge13);
        edges.add(edge14);
        edges.add(edge15);
        edges.add(edge16);
        graph.setEdges(edges);

        prim(graph,node1);




    }



}
